#!/bin/bash
# EbbFlow 本地测试脚本 - N=6, BatchSize=1000
set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 配置
N_NODES=6
EPOCHS=10
BATCH_SIZE=500
PROJECT_ROOT="/home/ebbflow"
LOG_DIR="${PROJECT_ROOT}/logs"
BINARY="${PROJECT_ROOT}/adkr_node"

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}==>${NC} ${GREEN}$1${NC}"
}

# 清理函数
cleanup() {
    log_step "清理进程和日志..."
    pkill -f adkr_node 2>/dev/null || true
    sleep 1
    rm -rf "$LOG_DIR"
    mkdir -p "$LOG_DIR"
    log_info "清理完成"
}

# 编译
build() {
    log_step "编译项目..."
    cd "$PROJECT_ROOT"
    go build -o adkr_node ./cmd/adkr_node/main.go
    if [ $? -ne 0 ]; then
        log_error "编译失败"
        exit 1
    fi
    log_info "编译成功"
}

# 检查配置
check_config() {
    log_step "检查配置文件..."
    
    # 检查 hosts.json
    if [ ! -f "$PROJECT_ROOT/hosts.json" ]; then
        log_error "hosts.json不存在，正在创建..."
        # 创建默认的hosts.json
        cat > "$PROJECT_ROOT/hosts.json" << EOF
[
    "localhost:11000",
    "localhost:11001",
    "localhost:11002",
    "localhost:11003",
    "localhost:11004",
    "localhost:11005"
]
EOF
        log_info "已创建默认hosts.json"
    fi
    
    log_info "配置文件验证通过: N=$N_NODES, Epochs=$EPOCHS, BatchSize=$BATCH_SIZE"
}

# 启动节点
start_nodes() {
    log_step "启动 $N_NODES 个节点..."
    
    for i in $(seq 0 $((N_NODES - 1))); do
        log_info "启动节点 $i..."
        nohup "$BINARY" -id=$i -epochs=$EPOCHS -fast > "$LOG_DIR/node_${i}.log" 2>&1 &
        echo $! > "$LOG_DIR/node_${i}.pid"
        sleep 0.5
    done
    
    sleep 2
    log_info "所有节点已启动"
}

# 检查节点状态
check_nodes() {
    log_step "检查节点状态..."
    
    local running=0
    for i in $(seq 0 $((N_NODES - 1))); do
        if [ -f "$LOG_DIR/node_${i}.pid" ]; then
            local pid=$(cat "$LOG_DIR/node_${i}.pid")
            if ps -p $pid > /dev/null 2>&1; then
                log_info "节点 $i (PID: $pid) 正在运行"
                running=$((running + 1))
            else
                log_error "节点 $i 已退出"
            fi
        fi
    done
    
    if [ $running -eq $N_NODES ]; then
        log_info "所有 $N_NODES 个节点都在运行"
        return 0
    else
        log_error "只有 $running/$N_NODES 个节点在运行"
        return 1
    fi
}

# 监控测试
monitor_test() {
    log_step "监控测试进度..."
    
    local timeout=300  # 5分钟超时
    local start_time=$(date +%s)
    local last_check=""
    
    while true; do
        local current_time=$(date +%s)
        local elapsed=$((current_time - start_time))
        
        # 检查超时
        if [ $elapsed -gt $timeout ]; then
            log_error "测试超时（${timeout}秒）"
            return 1
        fi
        
        # 检查是否有节点完成
        # 完成标志：ADKR-METRICS 且不包含 "error"，或者 Protocol finished successfully
        local completed_count=0
        for i in $(seq 0 $((N_NODES - 1))); do
            if [ -f "$LOG_DIR/node_${i}.log" ]; then
                # 检查是否有成功的ADKR-METRICS输出（不包含error）
                if grep -q "ADKR-METRICS" "$LOG_DIR/node_${i}.log" 2>/dev/null; then
                    # 检查是否包含error字段
                    if ! grep "ADKR-METRICS" "$LOG_DIR/node_${i}.log" | grep -q '"error"' 2>/dev/null; then
                        completed_count=$((completed_count + 1))
                    fi
                # 或者检查Protocol finished successfully
                elif grep -q "Protocol finished successfully" "$LOG_DIR/node_${i}.log" 2>/dev/null; then
                    completed_count=$((completed_count + 1))
                fi
            fi
        done
        
        # 显示进度
        local progress="[${elapsed}s] 完成: $completed_count/$N_NODES"
        if [ "$progress" != "$last_check" ]; then
            echo -ne "\r${progress}"
            last_check="$progress"
        fi
        
        # 如果所有节点都完成了
        if [ $completed_count -eq $N_NODES ]; then
            echo ""
            log_info "所有节点测试完成！"
            return 0
        fi
        
        # 检查是否有节点崩溃
        local running=0
        for i in $(seq 0 $((N_NODES - 1))); do
            if [ -f "$LOG_DIR/node_${i}.pid" ]; then
                local pid=$(cat "$LOG_DIR/node_${i}.pid")
                if ps -p $pid > /dev/null 2>&1; then
                    running=$((running + 1))
                fi
            fi
        done
        
        if [ $running -lt $N_NODES ]; then
            echo ""
            log_error "检测到节点崩溃！运行中: $running/$N_NODES"
            return 1
        fi
        
        sleep 2
    done
}

# 收集结果
collect_results() {
    log_step "收集测试结果..."
    
    echo ""
    echo "================================"
    echo "测试结果摘要"
    echo "================================"
    echo "配置: N=$N_NODES, Epochs=$EPOCHS, BatchSize=$BATCH_SIZE"
    echo ""
    
    for i in $(seq 0 $((N_NODES - 1))); do
        echo "--- 节点 $i ---"
        if [ -f "$LOG_DIR/node_${i}.log" ]; then
            # 提取关键指标
            if grep -q "ADKR-METRICS" "$LOG_DIR/node_${i}.log"; then
                grep "ADKR-METRICS" "$LOG_DIR/node_${i}.log" | tail -1 | grep -v '"error"' || echo "包含错误"
            elif grep -q "Protocol finished successfully" "$LOG_DIR/node_${i}.log"; then
                grep "Protocol finished successfully" "$LOG_DIR/node_${i}.log" | tail -1
            else
                echo "未完成"
            fi
        else
            echo "日志文件不存在"
        fi
        echo ""
    done
    
    # 检查是否有错误
    echo "================================"
    echo "错误检查"
    echo "================================"
    for i in $(seq 0 $((N_NODES - 1))); do
        if [ -f "$LOG_DIR/node_${i}.log" ]; then
            local error_count=$(grep -i "error\|panic\|fatal" "$LOG_DIR/node_${i}.log" | wc -l)
            if [ $error_count -gt 0 ]; then
                echo "节点 $i: $error_count 个错误"
                grep -i "error\|panic\|fatal" "$LOG_DIR/node_${i}.log" | head -5
            else
                echo "节点 $i: 无错误"
            fi
        fi
    done
    echo ""
}

# 显示日志文件位置
show_logs() {
    log_step "日志文件位置"
    echo "日志目录: $LOG_DIR"
    echo "查看日志: tail -f $LOG_DIR/node_0.log"
    echo ""
}

# 主函数
main() {
    echo ""
    echo "================================"
    echo "EbbFlow 本地测试"
    echo "配置: N=$N_NODES, Epochs=$EPOCHS, BatchSize=$BATCH_SIZE"
    echo "================================"
    echo ""
    
    # 清理
    cleanup
    
    # 编译
    build
    
    # 检查配置
    check_config
    
    # 启动节点
    start_nodes
    
    # 检查节点状态
    if ! check_nodes; then
        log_error "节点启动失败"
        show_logs
        exit 1
    fi
    
    # 监控测试
    if monitor_test; then
        # 收集结果
        collect_results
        
        # 清理进程
        log_step "清理进程..."
        pkill -f adkr_node 2>/dev/null || true
        
        log_info "测试完成！"
        show_logs
        exit 0
    else
        log_error "测试失败"
        show_logs
        
        # 显示最后几行日志
        echo ""
        log_step "最后的日志输出:"
        for i in $(seq 0 $((N_NODES - 1))); do
            echo "--- 节点 $i (最后20行) ---"
            tail -20 "$LOG_DIR/node_${i}.log" 2>/dev/null || echo "无日志"
            echo ""
        done
        
        exit 1
    fi
}

# 处理 Ctrl+C
trap 'echo ""; log_info "收到中断信号，清理进程..."; pkill -f adkr_node 2>/dev/null || true; exit 130' INT TERM

# 运行
main

